package org.zebra.web.websocket;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

import static org.zebra.web.websocket.WebSocketEmitterService.ATTRIBUTE_LOGIN_USER;

/**
 * 登录用户的 {@link HandshakeInterceptor} 实现类
 * <p>
 * 前端连接 websocket 时，会通过拼接 ?token={token} 到 ws:// 连接后
 *
 * @author zhanghongbin
 */
public class LoginInfoHandshakeInterceptor implements HandshakeInterceptor {

    private TokenHandler tokenHandler;

    public LoginInfoHandshakeInterceptor(TokenHandler tokenHandler) {
        this.tokenHandler = tokenHandler;
    }

    @Override
    public boolean beforeHandshake(
            ServerHttpRequest request,
            ServerHttpResponse response,
            WebSocketHandler wsHandler,
            Map<String, Object> attributes) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletServerHttpRequest = (ServletServerHttpRequest) request;
            LoginInfo loginInfo = this.tokenHandler.getLoginInfo(servletServerHttpRequest.getServletRequest());
            if (loginInfo == null) return false;
            attributes.put(ATTRIBUTE_LOGIN_USER, loginInfo);
        }
        return true;
    }

    @Override
    public void afterHandshake(
            ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {}
}
